import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np


def create_diffusion_process_visualization(save_path='diffusion_process_visualization.png'):
    """创建扩散过程可视化图，展示从噪声到检测结果的渐进过程"""

    # 设置画布 - 横向排列时间步
    fig, axes = plt.subplots(2, 5, figsize=(15, 7))

    # 定义时间步和对应的颜色强度
    timesteps = [800, 600, 400, 200, 0]

    # 创建模拟的藻类图像
    np.random.seed(42)

    # 图像尺寸
    img_width, img_height = 300, 200

    # 生成具有随机噪声的背景
    noise_level = 0.2
    background = np.random.normal(0.8, noise_level, (img_height, img_width, 3))
    background = np.clip(background, 0, 1)  # 限制在有效范围内

    # 添加一些模拟藻类的圆形
    num_algae = 25
    algae_positions = []

    for i in range(num_algae):
        x = np.random.randint(30, img_width - 30)
        y = np.random.randint(30, img_height - 30)
        r = np.random.randint(8, 15)
        color = (np.random.uniform(0.2, 0.4), np.random.uniform(0.5, 0.7), np.random.uniform(0.2, 0.4))
        algae_positions.append((x, y, r, color))

    # 创建藻类图像的副本用于所有时间步
    algae_img = background.copy()
    for x, y, r, color in algae_positions:
        # 创建藻类圆形
        for i in range(img_height):
            for j in range(img_width):
                # 如果在圆内
                if ((i - y) ** 2 + (j - x) ** 2) < r ** 2:
                    algae_img[i, j] = color

    # 生成每个时间步的图像和检测框
    for t_idx, timestep in enumerate(timesteps):
        # 计算当前时间步的噪声强度
        noise_strength = timestep / 800.0  # 从1到0
        signal_strength = 1 - noise_strength

        # 添加减少的噪声
        current_img = algae_img.copy() * signal_strength + np.random.normal(0.5, 0.3, algae_img.shape) * noise_strength
        current_img = np.clip(current_img, 0, 1)  # 限制在有效范围内

        # 显示当前时间步的图像
        axes[0, t_idx].imshow(current_img)
        axes[0, t_idx].set_title(f"t = {timestep}", fontsize=12)
        axes[0, t_idx].axis('off')

        # 创建对应的检测框 - 从噪声开始逐渐变得清晰
        box_img = np.ones((img_height, img_width, 3)) * 0.9  # 浅灰背景

        # 对于每个藻类位置添加一个检测框
        for x, y, r, color in algae_positions:
            # 基于时间步添加随机偏移 - 随着时间步减小，偏移变小
            offset_x = np.random.normal(0, 10 * noise_strength)
            offset_y = np.random.normal(0, 10 * noise_strength)

            # 添加大小抖动
            size_jitter = np.random.normal(0, 4 * noise_strength)

            # 计算边界框坐标
            box_left = max(0, int(x - r - 2 + offset_x))
            box_top = max(0, int(y - r - 2 + offset_y))
            box_right = min(img_width - 1, int(x + r + 2 + offset_x + size_jitter))
            box_bottom = min(img_height - 1, int(y + r + 2 + offset_y + size_jitter))

            # 检测框的可见度随时间步增加
            box_alpha = signal_strength * 0.8 + 0.2
            box_color = (1.0, 0.2, 0.2)  # 红色边框

            # 绘制矩形边框 - 顶部和底部线
            line_width = max(1, int(2 * signal_strength + 1))

            for w in range(line_width):
                for i in range(box_left, box_right + 1):
                    if 0 <= box_top + w < img_height and 0 <= i < img_width:
                        box_img[box_top + w, i] = [c * box_alpha + box_img[box_top + w, i, j] * (1 - box_alpha)
                                                   for j, c in enumerate(box_color)]
                    if 0 <= box_bottom - w < img_height and 0 <= i < img_width:
                        box_img[box_bottom - w, i] = [c * box_alpha + box_img[box_bottom - w, i, j] * (1 - box_alpha)
                                                      for j, c in enumerate(box_color)]

            # 绘制矩形边框 - 左侧和右侧线
            for w in range(line_width):
                for i in range(box_top, box_bottom + 1):
                    if 0 <= i < img_height and 0 <= box_left + w < img_width:
                        box_img[i, box_left + w] = [c * box_alpha + box_img[i, box_left + w, j] * (1 - box_alpha)
                                                    for j, c in enumerate(box_color)]
                    if 0 <= i < img_height and 0 <= box_right - w < img_width:
                        box_img[i, box_right - w] = [c * box_alpha + box_img[i, box_right - w, j] * (1 - box_alpha)
                                                     for j, c in enumerate(box_color)]

            # 如果接近最终时间步，添加置信度分数
            if timestep <= 200:
                # 置信度随时间增加
                conf = min(0.99, 0.7 + (200 - timestep) / 200 * 0.29 + np.random.normal(0, 0.05))
                conf_text = f"{conf:.2f}"
                font_size = int(9 * signal_strength + 6)

                # 添加文本标签到图像 - 简化版本，实际情况下会使用matplotlib的text功能
                text_x, text_y = box_left + 3, box_top - 3
                if box_top >= 15:  # 确保有足够空间显示文本
                    # 使用简单方式模拟文本，实际应用中可以使用PIL或其他方法
                    box_img[text_y - 10:text_y, text_x:text_x + 40] = box_color

        # 显示检测框图像
        axes[1, t_idx].imshow(box_img)
        axes[1, t_idx].set_title(f"Detection Boxes t = {timestep}", fontsize=12)
        axes[1, t_idx].axis('off')

    # 添加标题
    fig.suptitle("Diffusion Process Visualization for Microalgae Detection", fontsize=16, fontweight='bold')

    # 添加行标签
    fig.text(0.02, 0.75, "Input\nImage", fontsize=12, ha='center', fontweight='bold')
    fig.text(0.02, 0.25, "Detection\nBoxes", fontsize=12, ha='center', fontweight='bold')

    # 添加箭头指示扩散方向
    ax_arrow = fig.add_axes([0.25, 0.02, 0.5, 0.02])
    ax_arrow.arrow(0, 0.5, 1, 0, head_width=0.3, head_length=0.05, fc='black', ec='black')
    ax_arrow.text(0, 0, "Noise", fontsize=10, ha='center', va='center')
    ax_arrow.text(1, 0, "Clean", fontsize=10, ha='center', va='center')
    ax_arrow.axis('off')

    # 调整布局
    plt.tight_layout(rect=[0, 0.05, 1, 0.95])
    plt.subplots_adjust(wspace=0.05, hspace=0.2)

    # 保存图像
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.show()


create_diffusion_process_visualization()